# -*- coding:UTF-8 -*-
import logging
import pydantic
from pydantic import BaseModel
from typing import Any,List,Dict
from langchain.schema import Document
from langchain_core.prompts import PromptTemplate
from langchain.embeddings.huggingface import HuggingFaceEmbeddings
from langchain.utilities.google_serper import GoogleSerperAPIWrapper

from store.vectorstore import QdrantVectorStore
from data_loader.json_data_loader import LocalJSONLDataSetDatabase
from data_loader.file_data_loader import FileDataSetDatabase


LOG_FORMAT = "%(asctime)s - %(filename)s[line:%(lineno)d] - %(levelname)s: %(message)s"
logger = logging.getLogger()
logger.setLevel(logging.INFO)
logging.basicConfig(format=LOG_FORMAT)


SERVER_PORT=9886

SERVER_HOST="127.0.0.1"

QDRANT_URL='10.10.80.174:6333'

SERPAPI_API_KEY = "66a8b12acca8b3197fddd3e22907cec3f02cccc3"

# jsonl数据保存路径
LOCALDATA="data"

# 文件保存路径
FILES_DATA_DIR="data/files"

DIALOGUE_COLLECTION_NAME="validate_data"

RETRIEVAL_COLLECTION_NAME="opeba_retrieval_data"

RETRIEVAL_STORE_PATH="data/retriveal_db.json"

# 用于保存文件对应向量编号,{"文件名":[doc_id_01,doc_id_02,doc_id_03]}
FILE_AND_DOC_ID_STORGE=FileDataSetDatabase(store_path=RETRIEVAL_STORE_PATH)


# 向量库匹配阈值，阈值越高，匹配的内容越相似，如果为1，只有问题和向量库查找的内容完全相同时才会返回文本
SCORE_THRESHOLD=0.85

# OPENBA_3b_MODEL_PATH="/opt/data/private/liuteng/model/checkpoint-33"
OPENBA_3b_MODEL_PATH="/opt/data/private/liuteng/model/OpenBA-Chat-3B"
# OPENBA_3b_MODEL_PATH="/opt/data/private/liuteng/model/checkpoint-48"

BAICHUAN_4b_MODEL_PATH="/opt/data/private/liuteng/model/Baichuan2-7B-Chat"

EMBEDING_MODLE="/opt/data/private/liuteng/model/bge-large-zh"

class BaseResponse(BaseModel):
    code: int = pydantic.Field(200, description="API status code")
    msg: str = pydantic.Field("success", description="API status message")
    data: Any = pydantic.Field(None, description="API data")

    class Config:
        schema_extra = {
            "example": {
                "code": 200,
                "msg": "success",
            }
        }






TOOL_PROMPT='''下面给出五个例子，assistant根据reference参考资料与question提出的问题进行分析，如果reference参考资料能解答question的问题，assistant输出案例中给定的指定格式。
例1:
reference:不可以，需要到企业注册地的企业登记管理部门办理。
question:办理常熟的公司监事变更，可以在苏州工业园区办理吗？
assistant:{"content":"【reference中的内容可以回答question，所以reference与question之间有关系，不需要搜索工具】","is_need_search_tool":"no"}

例2:
reference:变更之后，董事会成员不得少于3人。董事名额的分配由合营各方参照出资比例协商确定。
user:合营企业变更董事会，对成员人数有什么要求吗？
assistant:{"content":"【reference中的内容可以回答question，所以reference与question之间有关系，不需要搜索工具】","is_need_search_tool":"no"}

例3:
reference:若为经办人或法定代表人领取，则带本人身份证即可；若为第三方（既不是法人也不是经办人）领取，则需携带加盖公司公章的介绍信以及办理人身份证领取。
user:苏州今天的天气怎么样？
assistant:{"content":"【reference中的内容无法回答question，reference与question之间没有关系，需要搜索工具】","is_need_search_tool":"yes"}

例4:
reference:我叫OpenBA-V2，是苏州大学自然语言研究室开发的
user:你叫什么名字？
assistant:{"content":"【reference中的内容可以回答question，所以reference与question之间有关系，不需要搜索工具】","is_need_search_tool":"no"}

例5:
reference:今天南昌天气正好，最适合散步了
user:一个股东可以设立公司吗？
assistant:{"content":"【reference中的内容无法回答question，reference与question之间没有关系，需要搜索工具】","is_need_search_tool":"yes"}

请根据上述案例格式输出json格式的assistant：
'''



SEARCH_PROMPT="请根据下面给定的参考内容：\n{content}\n请回答："



logger.info("正在加载向量数据库...")
dialogue_qdrant=QdrantVectorStore(
    url=QDRANT_URL,
    embeddings=HuggingFaceEmbeddings(model_name=EMBEDING_MODLE),
    collection_name=DIALOGUE_COLLECTION_NAME
)
retriveal_qdrant=QdrantVectorStore(
    url=QDRANT_URL,
    embeddings=HuggingFaceEmbeddings(model_name=EMBEDING_MODLE),
    collection_name=RETRIEVAL_COLLECTION_NAME
)
logger.info("向量数据库加载成功！")

logger.info("正在加载本地缓存数据库...")
localDataBase=LocalJSONLDataSetDatabase(id_key="id",dir_path=LOCALDATA)
logger.info("本地缓存数据库加载成功！")


logger.info("正在创建搜索工具")
search=GoogleSerperAPIWrapper(serper_api_key=SERPAPI_API_KEY ,k = 1,gl = "ch",hl= "zh-cn")
logger.info("正在创建搜索工具")

def create_tool_prompt(query):
    template="reference:{reference}\nuser:{query}\nassistant:"
    reference=search_in_vectorstore(query)
    prompt=PromptTemplate.from_template(template=template)
    return TOOL_PROMPT+prompt.invoke({"reference":reference,"query":query}).to_string()


def search_in_vectorstore(query:str):
    documents:List[Document]=dialogue_qdrant.search_document(text=query,k=1,score_threshold=SCORE_THRESHOLD)
    if documents is None or len(documents)<=0:
        return None
    content=""
    for document in documents:
        id=document.metadata["id"]
        json_data=localDataBase[id]
        content+=json_data["input"]+"\n"+json_data["output"]+"\n"
    print("【工具查询的内容：】",content)
    return content

def retriveal_docs_in_vecor(query:str):
    documents:List[Document]=retriveal_qdrant.search_document(text=query,k=4,score_threshold=0.7)
    if documents is None or len(documents)<=0:
        return None
    content=[]
    file_names=set()
    for document in documents:
        content.append(document.page_content)
        file_names.add(document.metadata["file_name"])
    text=query
    if len(content)>0:
        text=PromptTemplate.from_template(template=SEARCH_PROMPT).invoke({"content":'\n'.join(content)}).to_string()
    print("【工具查询的内容：】",text)
    print("【相似文件：】",file_names)
    return text,query,file_names


def search_in_net(query:str):
    try:
        res=search.run(query=query)
        print("【搜索引擎搜索的结果:】",res)
        return res
    except Exception as e:
        logger.error("搜索引擎调用失败")
    return None

def create_search_prompt(is_need_search_tool:False,query):
    # 如果问题相关，向向量库中查询是否包含对应内容
    content=None
    if not is_need_search_tool:
        content=search_in_vectorstore(query)
    
    # 如果向量库没有对应的问题和答案，则向搜索引擎查找
    if content is None:
        content=search_in_net(query)
    if content is None:
        return None
    return PromptTemplate.from_template(template=SEARCH_PROMPT).invoke({"content":content}).to_string() ,query


if __name__=="__main__":
    print(create_tool_prompt("在企业全链通综合服务平台中办理公司设立时，名称来源是什么？"))
    import pdb;pdb.set_trace()